package com.nbsaas.boot;

import com.baidu.aip.ocr.AipOcr;
import org.apache.pdfbox.pdmodel.PDDocument;
import org.apache.pdfbox.rendering.ImageType;
import org.apache.pdfbox.rendering.PDFRenderer;
import org.json.JSONArray;
import org.json.JSONObject;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.nio.file.Path;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;

/**
 * 针对扫描版 PDF（页面内容为图片），将每页渲染为图片并使用百度 OCR 提取文本。
 */
public class PdfImageOcrReader {

    private final AipOcr ocrClient;

    public PdfImageOcrReader() {
        // 使用配置中的百度 OCR 密钥初始化客户端
        this.ocrClient = new AipOcr(OcrConfig.BAIDU_APP_ID, OcrConfig.BAIDU_API_KEY, OcrConfig.BAIDU_SECRET_KEY);
        this.ocrClient.setConnectionTimeoutInMillis(OcrConfig.OCR_CONNECTION_TIMEOUT);
        this.ocrClient.setSocketTimeoutInMillis(OcrConfig.OCR_SOCKET_TIMEOUT);
    }

    /**
     * 对 PDF 进行 OCR，返回合并后的文本内容
     */
    public String extractText(Path pdfPath) throws IOException {
        List<String> pageTexts = new ArrayList<>();
        try (PDDocument document = PDDocument.load(pdfPath.toFile())) {
            PDFRenderer renderer = new PDFRenderer(document);
            int pageCount = document.getNumberOfPages();
            for (int i = 0; i < pageCount; i++) {
                BufferedImage image = renderer.renderImageWithDPI(i, 300, ImageType.RGB);
                byte[] imageBytes = bufferedImageToPngBytes(image);
                String text = ocrImageBytes(imageBytes);
                pageTexts.add(text);
            }
        }
        return String.join("\n\n", pageTexts).trim();
    }

    private byte[] bufferedImageToPngBytes(BufferedImage image) throws IOException {
        try (ByteArrayOutputStream baos = new ByteArrayOutputStream()) {
            ImageIO.write(image, "png", baos);
            return baos.toByteArray();
        }
    }

    private String ocrImageBytes(byte[] imageBytes) {
        HashMap<String, String> options = new HashMap<>();
        options.put("detect_direction", "true");
        options.put("probability", "true");
        JSONObject res = ocrClient.basicGeneral(imageBytes, options);
        if (res == null) {
            return "";
        }
        if (res.has("words_result")) {
            JSONArray arr = res.optJSONArray("words_result");
            if (arr != null) {
                StringBuilder sb = new StringBuilder();
                for (int i = 0; i < arr.length(); i++) {
                    JSONObject obj = arr.optJSONObject(i);
                    if (obj != null) {
                        String w = obj.optString("words", "");
                        if (!w.isEmpty()) {
                            sb.append(w).append("\n");
                        }
                    }
                }
                return sb.toString().trim();
            }
        }
        return res.optString("error_msg", "");
    }
}